import React, { useState, useEffect } from 'react';
import { Layout, Button, Spin, Input, Select, Progress, Descriptions, Divider, ConfigProvider, Result, Upload, message, Space } from 'antd';
import { UploadOutlined, ArrowLeftOutlined } from '@ant-design/icons';
import { useNavigate, useLocation } from 'react-router-dom';
import axios from 'axios';

const { Option } = Select;
const { Header, Content } = Layout;

const UploadPage = () => {
  const [currentStep, setCurrentStep] = useState(0);
  const [loading, setLoading] = useState(false);
  const [progress, setProgress] = useState(0);
  const [fileList, setFileList] = useState([]);
  const [uploading, setUploading] = useState(false);
  const [modelInfo, setModelInfo] = useState({
    model_id: '',
    name: '',
    location: 'Remote',
    type: '',
    attribute: '',
    raw:'pth',
    input_size:'',
    final:'',
    state: 'uploading'
  });
  const [validationInfo, setValidationInfo] = useState({});
  const [compressionInfo, setCompressionInfo] = useState({});
  const [testInfo, setTestInfo] = useState({});
  const [canProceed, setCanProceed] = useState(false);
  const navigate = useNavigate();
  const location = useLocation();

  useEffect(() => {
    if (location.state && location.state.state === 'ready') {
      setModelInfo({
        model_id: location.state.modelId || '',
        name: location.state.name || '',
        location: location.state.location || '',
        type: location.state.type || '',
        attribute: location.state.attribute || '',
        input_size: location.state.input_size || '', // Added input_size
        final: location.state.final || '', // Added final
        state: 'uploading'
      });

      setCurrentStep(0);
    } else if (location.state) {
      setModelInfo({
        model_id: location.state.modelId || '',
        name: location.state.name || '',
        location: location.state.location || '',
        type: location.state.type || '',
        attribute: location.state.attribute || '',
        input_size: location.state.input_size || '', // Added input_size
        final: location.state.final || '', // Added final
        state: location.state.state
      });

      const stepMap = {
        uploading: 0,
        validate: 1,
        compress: 2,
        evaluate: 3,
        ready: 4
      };
      setCurrentStep(stepMap[location.state.state] || 0);
    }
  }, [location.state]);

  const steps = [
    {
      title: '开始上传',
      description: '开始上传模型',
    },
    {
      title: '模型验证',
      description: '正在进行模型验证',
    },
    {
      title: '模型压缩',
      description: '正在压缩模型',
    },
    {
      title: '模型测试',
      description: '正在测试模型',
    },
    {
      title: '上传结束',
      description: '模型上传成功',
    },
  ];

  useEffect(() => {
    setCurrentStep(0);
    setLoading(false);
    setProgress(0);
  }, []);

  const handleNextStep = () => {
    if (currentStep < steps.length - 1) {
      setLoading(true);
      setTimeout(() => {
        setCurrentStep(currentStep + 1);
        setLoading(false);
        setProgress(0);
        setCanProceed(false); // 每次进入新步骤时重置canProceed
      }, 1000);
    }
  };

  const handleUpload = async () => {
    setUploading(true);

    // 模拟上传进度
    let currentProgress = 0;
    const progressInterval = setInterval(() => {
      currentProgress += 20;
      setProgress(currentProgress);

      if (currentProgress >= 100) {
        clearInterval(progressInterval);
        // 5秒后显示上传成功
        setTimeout(() => {
          setUploading(false);
          message.success('🎉 模型上传成功！');
          // 重置表单
          handleUploadAgain();
        }, 1000); // 进度条完成后再等1秒
      }
    }, 1000); // 每秒增加20%，总共5秒
  };

  const handleProgress = () => {
    setProgress(0);
    const interval = setInterval(() => {
      setProgress((prev) => {
        if (prev >= 100) {
          clearInterval(interval);
          return 100;
        }
        return Math.round(prev + 2.33); // 3秒内从0%到100%，并四舍五入为整数
      });
    }, 100);
  };

  useEffect(() => {
    if (currentStep > 0 && currentStep < 4) {
      handleProgress();
    }
  }, [currentStep]);

  const handleUploadAgain = () => {
    setCurrentStep(0);
    setLoading(false);
    setProgress(0);
    setUploading(false);
    setFileList([]);
    setModelInfo({
      model_id: '',
      name: '',
      location: 'Remote',
      type: '',
      attribute: '',
      input_size: '',
      final: '',
      state: 'uploading'
    });
  };

  const handleBack = () => {
    navigate('/');
    window.location.reload();
  };

  const uploadProps = {
    onRemove: (file) => {
      const index = fileList.indexOf(file);
      const newFileList = fileList.slice();
      newFileList.splice(index, 1);
      setFileList(newFileList);
    },
    beforeUpload: (file) => {
      setFileList([file]); // 只保留一个文件
      return false;
    },
    fileList,
    maxCount: 1
  };

  const handleValidation = async () => {
    try {
      const response = await axios.get(`http://39.106.250.42:5172/model_validate/${modelInfo.model_id}`);
      setValidationInfo(response.data);
      setCanProceed(true); // 设置canProceed为true
    } catch (error) {
      message.error('Model validation failed.');
    }
  };

  const handleCompression = async () => {
    try {
      const response = await axios.get(`http://39.106.250.42:5172/model_compress/${modelInfo.model_id}`);
      setCompressionInfo(response.data);
      setCanProceed(true); // 设置canProceed为true
    } catch (error) {
      message.error('Model compression failed.');
    }
  };

  const handleTest = async () => {
    try {
      const response = await axios.get(`http://39.106.250.42:5172/model_test/${modelInfo.model_id}`);
      setTestInfo(response.data);
      setCanProceed(true); // 设置canProceed为true
    } catch (error) {
      message.error('Model test failed.');
    }
  };

  useEffect(() => {
    if (currentStep === 1) {
      handleValidation();
    } else if (currentStep === 2) {
      handleCompression();
    } else if (currentStep === 3) {
      handleTest();
    }
  }, [currentStep]);

  return (
    <Layout className='ant-layout1' style={{margin: '0', height: '95%', backgroundColor: 'transparent', border: 'none'}}>
      <Layout className='ant-layout2' style={{backgroundColor:'transparent'}}>
        <Content className='ant-layout-content' style={{ padding: '10px', display: 'flex', flexDirection: 'row', justifyContent: 'flex-start', alignItems: 'flex-start', position: 'relative' }}>
          <div style={{ width: '100%', maxWidth: '600px', paddingLeft: '20px' }}>
          {currentStep === 0 && (
            <>
              <p className='input-label'>模型ID: </p>
              <Input
                placeholder="输入模型ID"
                value={modelInfo.model_id}
                onChange={(e) => setModelInfo({ ...modelInfo, model_id: e.target.value })}
                style={{ marginBottom: 10, width: '100%' }}
                disabled={uploading}
              />
              <p className='input-label'>名称: </p>
              <Input
                placeholder="输入模型名称"
                value={modelInfo.name}
                onChange={(e) => setModelInfo({ ...modelInfo, name: e.target.value })}
                style={{ marginBottom: 10, width: '100%' }}
                disabled={uploading}
              />
              <p className='input-label'>位置: </p>
              <Input
                placeholder="Remote"
                value={modelInfo.location}
                onChange={(e) => setModelInfo({ ...modelInfo, location: e.target.value })}
                style={{ marginBottom: 10, width: '100%' }}
                defaultValue={'Remote'}
                disabled
              />
              <p className='input-label'>选择模型类型</p>
              <Select
                placeholder="选择模型类型"
                value={modelInfo.type}
                onChange={(value) => setModelInfo({ ...modelInfo, type: value })}
                style={{ width: '100%', marginBottom: 10 }}
                disabled={uploading}
              >
                <Option value="图像分类">图像分类</Option>
                <Option value="雷达识别">雷达识别</Option>
                <Option value="变化检测">变化检测</Option>
                <Option value="云检测">云检测</Option>
                <Option value="导航识别">导航识别</Option>
                <Option value="目标检测">目标检测</Option>
                <Option value="图像分割">图像分割</Option>
              </Select>
              <p className='input-label'>描述: </p>
              <Input
                placeholder="输入描述"
                value={modelInfo.attribute}
                onChange={(e) => setModelInfo({ ...modelInfo, attribute: e.target.value })}
                style={{ marginBottom: 10, width: '100%' }}
                disabled={uploading}
              />
              <p className='input-label'>输入大小: </p>
              <Input
                placeholder="输入大小"
                value={modelInfo.input_size}
                onChange={(e) => setModelInfo({ ...modelInfo, input_size: e.target.value })}
                style={{ marginBottom: 10, width: '100%' }}
                disabled={uploading}
              />
              <p className='input-label'>选择最终格式</p>
              <Select
                placeholder="选择最终格式"
                value={modelInfo.final}
                onChange={(value) => setModelInfo({ ...modelInfo, final: value })}
                style={{ width: '100%', marginBottom: 10 }}
                disabled={uploading}
              >
                <Option value="onnx">onnx</Option>
                <Option value="mnn">mnn</Option>
              </Select>
              <Divider/>

              {/* 上传进度显示 */}
              {uploading && (
                <div style={{ marginBottom: 20 }}>
                  <p>正在上传模型文件...</p>
                  <Progress
                    percent={progress}
                    status="active"
                    strokeColor={{
                      '0%': '#108ee9',
                      '100%': '#87d068',
                    }}
                  />
                </div>
              )}

              <Space>
                <Upload {...uploadProps} disabled={uploading}>
                  <Button disabled={uploading}>选择文件</Button>
                </Upload>
                <Divider type='vertical'/>
                <Button
                  type="primary"
                  onClick={handleUpload}
                  disabled={fileList.length === 0 || !modelInfo.model_id || !modelInfo.name || !modelInfo.location || !modelInfo.type || !modelInfo.attribute || !modelInfo.input_size || !modelInfo.final}
                  loading={uploading}
                  icon={<UploadOutlined />}
                >
                  {uploading ? '上传中...' : '开始上传'}
                </Button>
              </Space>
            </>
          )}

            {currentStep === 1 && (
              <>
                <h3>模型验证</h3>
                <Progress percent={progress} style={{ marginBottom: 20 }} />
                <Divider/>
                <Descriptions title="模型验证信息">
                  <Descriptions.Item label="验证结果">{validationInfo.result}</Descriptions.Item>
                  <Descriptions.Item label="验证时间">{validationInfo.time}</Descriptions.Item>
                </Descriptions>
              </>
            )}
            {currentStep === 2 && (
              <>
                <h3>模型压缩</h3>
                <Progress percent={progress} style={{ marginBottom: 20 }} />
                <Divider/>
                <Descriptions title="压缩效果">
                  <Descriptions.Item label="未压缩体积">{compressionInfo.originalSize}</Descriptions.Item>
                  <Descriptions.Item label="未压缩正确率">{compressionInfo.originalAccuracy}</Descriptions.Item>
                  <Descriptions.Item label="压缩体积">{compressionInfo.compressedSize}</Descriptions.Item>
                  <Descriptions.Item label="压缩正确率">{compressionInfo.compressedAccuracy}</Descriptions.Item>
                </Descriptions>
              </>
            )}
            {currentStep === 3 && (
              <>
                <h3>模型测试</h3>
                <Progress percent={progress} style={{ marginBottom: 20 }} />
                <Divider/>
                <Descriptions title="测试结果">
                  <Descriptions.Item label="测试结果">{testInfo.result}</Descriptions.Item>
                  <Descriptions.Item label="测试时间">{testInfo.time}</Descriptions.Item>
                </Descriptions>
              </>
            )}
            {currentStep === 4 && (
              <Result
                status="success"
                title="Successfully Upload New Model!"
              />
            )}
          </div>
          <div style={{ position: 'absolute', bottom: '15vh', textAlign: 'center', width: '100%' }}>
            {currentStep === 0}
            {currentStep > 0 && currentStep < 4 && (
              <Button
                type="primary"
                onClick={handleNextStep}
                disabled={!canProceed || progress < 100} // 修改按钮的disabled属性
              >
                Next Step
              </Button>
            )}
            {currentStep === 4 && (
              <Button
                type="primary"
                onClick={handleUploadAgain}
              >
                Upload Again
              </Button>
            )}
          </div>
        </Content>
      </Layout>
    </Layout>
  );
};

export default UploadPage;